from abc import ABC
from typing import List

from langchain.chains.base import Chain
from langchain_core.language_models import BaseChatModel
from pydantic import Extra


class BaseChain(Chain, ABC):
    llm: BaseChatModel
    output_key: str = "answer"

    class Config:
        extra = Extra.forbid
        arbitrary_types_allowed = True

    @property
    def input_keys(self) -> List[str]:
        """Expect input key.
        :meta private:
        """
        return ["question", "chat_history"]

    @property
    def output_keys(self) -> List[str]:
        """Return output key.
        :meta private:
        """
        return [self.output_key]
